#!/usr/bin/env bash
# yitong.bai@qq.com
# Library for ssh copy id functions
# 还不太好使呢 ! ! !

# Load functions
cd $(dirname $(readlink -f $0)) && source liblog.sh

# Functions

########################
# Install Rpm Requirements
# Arguments:
#   None
# Returns:
#   None
#########################
install_rpm_requirements_ssh_copy_id() {
  info "*** Installing rpm requirements ***"
  yum -y install expect > /dev/null 
}

########################
# Generate hosts list from ansible host inventory file
# Arguments:
#   ansible host inventory file path
# Returns:
#   None
#########################
function host_list_gen() {
  ansible_inventory_host=$1
  if [[ -n "${ansible_inventory_host}" ]]
  then
    grep -E -o "([0-9]{1,3}[\.]){3}[0-9]{1,3}" "${ansible_inventory_host}" | sort | uniq > /tmp/host_list.txt
  else
    info "*** No avalible ansible inventory host found, use default /tmp/host_list.txt instead ***"
  fi
}


########################
# Generate public/private rsa key pair
# Arguments:
#   None
# Returns:
#   None
#########################
function ssh_key_gen() {
  info "*** Generating public/private rsa key pair ***"
  expect << __EOF
  set timeout 30
  spawn ssh-keygen
  expect {
    ".ssh/id_rsa):" { send "\r" ; exp_continue }
    "Overwrite (y/n)?" { send "n\r" ; exp_continue }
    "(empty for no passphrase):" { send "\r" ; exp_continue }
    "again:" { send "\r" ; exp_continue }
  }
  interact
__EOF
  sleep 1
}

########################
# Auto ssh copy id
# Arguments:
#   username and password
# Returns:
#   None
#########################
function auto_ssh_copy_id() {
  username=$1
  password=$2

  if [[ -z "${password}" ]]
  then 
    username="root"
    password=$1
  fi

  for server in $(cat /tmp/host_list.txt )
  do
    info "*** copy id to server ${server} ***"
    expect << __EOF
    set timeout 30
    spawn ssh-copy-id ${username}@${server}
    expect {
      "(yes/no)?" { send "yes\r" ; exp_continue }
      "password:" { send "${password}\r" ; exp_continue }
    }
    interact
__EOF
  done
}

########################
# Auto ssh copy id with sudo user, 
# Arguments:
#   username and password
# Returns:
#   None
#########################
function auto_ssh_copy_id_with_sudo(){
  #因为有密码的sudoers不允许在没有tty终端的情况下sudo。例如ssh连接进行sudo，会报错sudo: no tty present and no askpass program specified。
  #可以修改/etc/sudoers，userABC  ALL=(ALL) NOPASSWD，不要密码验证。或者注释sudoers配置文件中 #Defaults requiretty。

  username=$1
  password=$2

  user_home_dir=$3
  [[ -z "${user_home_dir}" ]] && user_home_dir=/home/"${username}"
  
  for server in $(cat /tmp/host_list.txt )
  do
    info "*** scp public key to ${server} ***"
    expect << __EOF
      set timeout 30
      spawn scp ${user_home_dir}/.ssh/id_rsa.pub ${username}@${server}:${user_home_dir}
      expect {
        "(yes/no)?" { send "yes\r" ; exp_continue }
        "password:" { send "${password}\r" ; exp_continue }
      }

      spawn ssh -t ${username}@${server}
          sudo mkdir -p /root/.ssh \
      &&  sudo chmod 644 /root/.ssh \
      &&  sudo cp ${user_home_dir}/id_rsa.pub /root/.ssh/authorized_keys \
      &&  sudo chown root.root /root/.ssh/authorized_keys \
      &&  sudo chmod 600 /root/.ssh/authorized_keys
      "
      expect {
        "(yes/no)?" { send "yes\r" ; exp_continue }
        "password:" { send "${password}\r" ; exp_continue }
        "password for ${username}:" { send "${password}\r" ; exp_continue }
      }
      interact
__EOF
  done

  info "*** copy private key to root@localhost ***"
  expect << __EOF
    set timeout 10
    spawn sudo cp ${user_home_dir}/.ssh/id_rsa ${user_home_dir}/.ssh/known_hosts /root/.ssh
    expect {
        "password for ${username}:" { send "${password}\r" ; exp_continue }
    }
    interact
__EOF
}
